﻿/*
 * CUMediaWriter.cpp
 *
 *  Created on: 2020年7月18日
 *      Author: terry
 */

#include "CUMediaWriter.h"
#include "CLog.h"
#include "Path.h"
#include "CodecIDConvert.h"
#include "AacHelper.h"
#include "UMediaFileFormat.h"


CUMediaWriter::CUMediaWriter(int fileFormat):
    m_fileFormat(fileFormat),
    m_format(),
    m_fileName(),
    m_changeFile(),
    m_fmtCtx(),
    m_videoIndex(-1),
    m_audioIndex(-1),
	m_subtitleIndex(-1),
    m_videoStartTime(AV_NOPTS_VALUE),
    m_videoCurrentTime(),
    m_audioStartTime(AV_NOPTS_VALUE),
	m_subtitleStartTime(AV_NOPTS_VALUE),
    m_audioCurrentTime(),
	m_opened()
{
}

CUMediaWriter::~CUMediaWriter()
{
    close();
}


bool CUMediaWriter::open(const char* filename, const UMediaFormat& fmt)
{
    if (!filename || strlen(filename) == 0)
    {
        return false;
    }

    m_fileName = filename;
    m_format = fmt;

	m_opened = openFile();

	return m_opened;
}

void CUMediaWriter::close()
{
    closeFile();
}

bool CUMediaWriter::isOpen()
{
    return isFileOpen();
}

bool CUMediaWriter::getFormat(UMediaFormat& fmt)
{
    fmt = m_format;
	return true;
}

bool CUMediaWriter::write(UMediaPacket& pkt)
{
    if (pkt.isVideo())
    {
        writeVideo(*pkt.get());
    }
    else if (pkt.isAudio())
    {
		if (m_audioIndex >= 0)
		{
			writeAudio(*pkt.get());
		}
    }
	else if (pkt.type() == UMEDIA_TYPE_SUBTITLE)
	{
		writeSubtitle(*pkt.get());
	}
	else
	{
		writeSubtitle(*pkt.get());
	}

	return true;
}

/**
 * 设置输出文件路径. 可用于切换录像
 * @param filepath	文件路径
 * @return true 表示成功
 */
bool CUMediaWriter::setFile(const char* filepath)
{
    if (filepath == NULL)
    {
        m_fileName.clear();
    }
    else
    {
        m_fileName = filepath;
    }

    {
        std::lock_guard< std::mutex >   lock(m_mutex);
        if (isFileOpen())
        {
            m_changeFile = true;
        }
    }

    return true;
}

int CUMediaWriter::getFileFormat()
{
	return m_fileFormat;
}

bool CUMediaWriter::openFile()
{
    int rc = 0;

    rc = avformat_alloc_output_context2(&m_fmtCtx, NULL, NULL, m_fileName.c_str());
    if (!m_fmtCtx)
    {
        std::string fileDir = comn::Path::dirname(m_fileName);
        if (!comn::Path::exist(fileDir))
        {
            if (comn::Path::createDirectories(fileDir))
            {
                CLog::info("create directory: %s\n", fileDir.c_str());

                rc = avformat_alloc_output_context2(&m_fmtCtx, NULL, NULL, m_fileName.c_str());
            }
            else
            {
                CLog::error("failed to create dir: %s\n", fileDir.c_str());
            }
        }
    }

    if (!m_fmtCtx)
    {
        CLog::warning("can not alloc avformat. %s\n", m_fileName.c_str());
        return false;
    }

    if (m_format.hasVideo())
    {
        AVCodecID codecID = CodecIDConvert::toAVCodecID(m_format.m_codec);
        AVCodec* codec = avcodec_find_decoder(codecID);
        if (codec != NULL)
        {
            //CLog::debug("video codec: %s\n", codec->name);

			AVStream *out_stream = avformat_new_stream(m_fmtCtx, NULL);

#if LIBAVCODEC_VERSION_MAJOR < 58
            
            AVCodecContext* codecCtx = out_stream->codec;

            avcodec_get_context_defaults3(codecCtx, codec);

            if (m_format.m_width <= 0)
            {
                codecCtx->width = 1920;
                codecCtx->height = 1080;
            }
            else
            {
                codecCtx->width = m_format.m_width;
                codecCtx->height = m_format.m_height;
                codecCtx->profile = m_format.m_profile;
            }

            if (m_format.m_framerate > 0)
            {
                codecCtx->framerate = av_make_q(m_format.m_framerate, 1);

                out_stream->r_frame_rate = av_make_q(m_format.m_framerate, 1);
                out_stream->avg_frame_rate = av_make_q(m_format.m_framerate, 1);
            }

            codecCtx->time_base = av_make_q(1, 90000);

            out_stream->time_base = codecCtx->time_base;

            // extra_data
            if (m_format.getVideoPropSize() > 0)
            {
                codecCtx->extradata = (uint8_t*)av_memdup(m_format.getVideoProp(), m_format.getVideoPropSize());
                codecCtx->extradata_size = m_format.getVideoPropSize();
            }
#else

			out_stream->time_base = av_make_q(1, AV_TIME_BASE);

			AVCodecParameters* codecpar = out_stream->codecpar;
			codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
			codecpar->codec_id = codecID;
			codecpar->width = m_format.m_width > 0 ? m_format.m_width : 1920;
			codecpar->height = m_format.m_height > 0 ? m_format.m_height : 1080;
			codecpar->profile = m_format.m_profile > 0 ? m_format.m_profile : 0;


			if (m_format.getVideoPropSize() > 0)
			{
				codecpar->extradata_size = m_format.getVideoPropSize();
				codecpar->extradata = (uint8_t*)av_memdup(m_format.getVideoProp(), m_format.getVideoPropSize());
			}

#endif //

			if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
			{
				out_stream->codec->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
			}

			m_videoIndex = 0;
        }
        else
        {
            CLog::warning("no such video codec %d\n", codecID);
        }
    }

    if (m_format.hasAudio())
    {
        AVCodecID codecID = CodecIDConvert::toAVCodecID(m_format.m_audioCodec);

		bool available = checkAudioSupport(m_fileFormat, m_fmtCtx, codecID);

        AVCodec* codec = avcodec_find_decoder(codecID);
        if (codec != NULL && available)
        {
            //CLog::debug("audio codec: %s\n", codec->name);

			AVStream *out_stream = avformat_new_stream(m_fmtCtx, codec);

#if LIBAVCODEC_VERSION_MAJOR < 58

            AVCodecContext* codecCtx = out_stream->codec;

            avcodec_get_context_defaults3(codecCtx, codec);


            codecCtx->channels = m_format.m_channels;
            codecCtx->channel_layout = av_get_default_channel_layout(m_format.m_channels);
            codecCtx->sample_rate = m_format.m_sampleRate;
            codecCtx->profile = m_format.m_audioProfile;
            codecCtx->time_base.num = 1;
            codecCtx->time_base.den = codecCtx->sample_rate;

            out_stream->time_base = codecCtx->time_base;

            if (m_format.getAudioConfigSize() > 0)
            {
                codecCtx->extradata = (uint8_t*)av_memdup(m_format.getAudioConfig(), m_format.getAudioConfigSize());
                codecCtx->extradata_size = m_format.getAudioConfigSize();

                //std::string extra = comn::StringCast::toHexString(codecCtx->extradata, //codecCtx->extradata_size);
                //CLog::debug("extra: %s\n", extra.c_str());
            }

#else
			out_stream->time_base = av_make_q(1, m_format.m_sampleRate);

			AVCodecParameters* codecpar = out_stream->codecpar;
			codecpar->codec_type = AVMEDIA_TYPE_AUDIO;
			codecpar->codec_id = codecID;
			codecpar->channels = m_format.m_channels;
			codecpar->sample_rate = m_format.m_sampleRate;
			codecpar->bit_rate = m_format.m_audioBitrate;

			codecpar->bits_per_raw_sample = 16;
			codecpar->bits_per_coded_sample = codecpar->bits_per_raw_sample;

			if (codecID == AV_CODEC_ID_AAC)
			{
				codecpar->frame_size = m_format.m_frameSize > 0 ? m_format.m_frameSize : 1024;
			}
			else
			{
				codecpar->frame_size = m_format.m_frameSize;
			}
			//codecpar->block_align = codecpar->frame_size;
			
			codecpar->channel_layout = av_get_default_channel_layout(m_format.m_channels);

			if (m_format.getAudioConfigSize() > 0)
			{
				codecpar->extradata_size = m_format.getAudioConfigSize();
				codecpar->extradata = (uint8_t*)av_memdup(m_format.getAudioConfig(), m_format.getAudioConfigSize());
			}

#endif //
			if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
			{
				out_stream->codec->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
			}

			m_audioIndex = m_videoIndex + 1;
        }
    }

	if (m_format.m_subtitle > 0)
	{
		AVCodecID codecID = AV_CODEC_ID_MOV_TEXT;// CodecIDConvert::toAVCodecID(m_format.m_subtitle);
		AVCodec* codec = avcodec_find_encoder(codecID);
		if (codec != NULL)
		{
			AVStream *out_stream = avformat_new_stream(m_fmtCtx, codec);

#if LIBAVCODEC_VERSION_MAJOR < 58
            AVCodecContext* codecCtx = out_stream->codec;

            avcodec_get_context_defaults3(codecCtx, codec);

            codecCtx->time_base = av_make_q(1, AV_TIME_BASE);

            codecCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
            codecCtx->codec_type = AVMEDIA_TYPE_SUBTITLE;
            codecCtx->codec_id = codecID;
#else
			out_stream->time_base = av_make_q(1, AV_TIME_BASE);

			AVCodecParameters* codecpar = out_stream->codecpar;
			codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
			codecpar->codec_id = codecID;
#endif //

			m_subtitleIndex = m_audioIndex + 1;
		}
	}

    rc = avio_open(&m_fmtCtx->pb, m_fileName.c_str(), AVIO_FLAG_WRITE);
    if (m_fmtCtx->pb == NULL)
    {
        CLog::error("failed to open file %s\n", m_fileName.c_str());

		closeFile();

        return false;
    }

    rc = avformat_write_header(m_fmtCtx, NULL);
    if (rc != 0)
    {
        CLog::error("failed to write header %s\n", m_fileName.c_str());

		closeFile();
        
        return false;
    }

    return true;
}

void CUMediaWriter::closeFile()
{
    if (m_fmtCtx)
    {
        if (m_fmtCtx->pb)
        {
			if (m_opened)
			{
				av_write_trailer(m_fmtCtx);
			}

            avio_close(m_fmtCtx->pb);
        }

        for (unsigned int i = 0; i < m_fmtCtx->nb_streams; i ++)
        {
            AVStream* stream = m_fmtCtx->streams[i];
            if (stream && stream->codec)
            {
                avcodec_close(stream->codec);
            }
        }

        avformat_free_context(m_fmtCtx);
        m_fmtCtx = NULL;
    }

    m_changeFile = false;

    m_videoIndex = -1;
    m_audioIndex = -1;
	m_subtitleIndex = -1;
    m_videoStartTime = AV_NOPTS_VALUE;
    m_audioStartTime = AV_NOPTS_VALUE;
	m_subtitleStartTime = AV_NOPTS_VALUE;
    m_videoCurrentTime = 0;
    m_audioCurrentTime = 0;

	m_opened = false;
}


bool CUMediaWriter::isFileOpen()
{
    return (m_fmtCtx != NULL && m_opened);
}


void CUMediaWriter::writeVideo(AVPacket& inPkt)
{
    if (m_videoIndex < 0)
    {
        CLog::warning("file is not ready. video index < 0. size:%d, flags:%d\n", inPkt.size, inPkt.flags);
        return;
    }

    if (m_videoStartTime == AV_NOPTS_VALUE)
    {
        m_videoStartTime = inPkt.dts;
        m_videoCurrentTime = m_videoStartTime;
    }

    int64_t dts = inPkt.dts - m_videoStartTime;

    int64_t duration = inPkt.dts - m_videoCurrentTime;
    if (duration < 0)
    {
        duration = 0;
    }

    m_videoCurrentTime = inPkt.dts;

    AVPacket pkt;
    av_init_packet(&pkt);

    av_packet_ref(&pkt, &inPkt);

    pkt.stream_index = m_videoIndex;
    pkt.pts = dts;
    pkt.dts = dts;
    pkt.duration = duration;

    AVRational src = av_make_q(1, AV_TIME_BASE);
    AVRational dst = m_fmtCtx->streams[m_videoIndex]->time_base;
    av_packet_rescale_ts(&pkt, src, dst);

  //  CLog::debug("video packet. size:%d, duration:%d, dts:%lld\n",
        //pkt.size, (int)pkt.duration, pkt.dts);

    int error = av_interleaved_write_frame(m_fmtCtx, &pkt);
    if (error < 0)
    {
        CLog::warning("failed to write video frame. error: %d, in pkt size: %d, dts: %lld\n", error, inPkt.size, inPkt.dts);
    }

    av_packet_unref(&pkt);
}

void CUMediaWriter::writeAudio(AVPacket& inPkt)
{
    //CLog::debug("CUMediaWriter::writeAudio. size:%d, dts:%lld\n",  inPkt.size, inPkt.dts);

    if (m_audioStartTime == AV_NOPTS_VALUE)
    {
        m_audioStartTime = inPkt.pts;
    }

    int64_t pts = inPkt.pts - m_audioStartTime;

    AVPacket pkt;
    av_init_packet(&pkt);

	if (m_fileFormat == UMEDIA_FILE_FLV && AacHelper::isAdts(inPkt.data, inPkt.size))
	{
		pkt.data = inPkt.data + AacHelper::ADTS_HEADER_SIZE;
		pkt.size = inPkt.size - AacHelper::ADTS_HEADER_SIZE;

		pkt.stream_index = m_audioIndex;
		pkt.pts = pts;
		pkt.dts = pts;

		doWriteAudio(&pkt, AV_TIME_BASE);
	}
	else
	{
		av_packet_ref(&pkt, &inPkt);
		pkt.stream_index = m_audioIndex;
		pkt.pts = pts;
		pkt.dts = pts;

		doWriteAudio(&pkt, AV_TIME_BASE);

		av_packet_unref(&pkt);
	}

}

void CUMediaWriter::doWriteAudio(AVPacket* pkt, int rate)
{
    //CLog::debug("Audio packet. size:%d, dts:%lld\n",  pkt->size, pkt->dts);

    if(m_fmtCtx)
    {
        AVRational src = av_make_q(1, rate);
        AVRational dst = m_fmtCtx->streams[m_audioIndex]->time_base;
        av_packet_rescale_ts(pkt, src, dst);

        int64_t dts = pkt->dts;
        int error = av_interleaved_write_frame(m_fmtCtx, pkt);
        if (error < 0)
        {
            CLog::warning("failed to write audio frame. error: %d, dts: %lld\n", error, dts);
        }
    }

    //CLog::debug("Audio packet end.\n");
}

void CUMediaWriter::writeSubtitle(AVPacket& inPkt)
{
	if (m_subtitleStartTime == AV_NOPTS_VALUE)
	{
		m_subtitleStartTime = inPkt.pts;
	}

	int64_t pts = inPkt.pts - m_subtitleStartTime;

	AVPacket pkt;
	av_init_packet(&pkt);

	av_packet_ref(&pkt, &inPkt);
	pkt.stream_index = m_subtitleIndex;
	pkt.pts = pts;
	pkt.dts = pts;

	int rate = AV_TIME_BASE;

	AVRational src = av_make_q(1, rate);
	AVRational dst = m_fmtCtx->streams[m_audioIndex]->time_base;
	av_packet_rescale_ts(&pkt, src, dst);

	int64_t dts = pkt.dts;
	int error = av_interleaved_write_frame(m_fmtCtx, &pkt);
	if (error < 0)
	{
		CLog::warning("failed to write audio frame. error: %d, dts: %lld\n", error, dts);
	}

	av_packet_unref(&pkt);
}

bool CUMediaWriter::checkAudioSupport(int fileFormat, AVFormatContext* fmtCtx, AVCodecID codecID)
{
	bool got = true;
	if (fileFormat == UMEDIA_FILE_MP4)
	{
		if (codecID == AV_CODEC_ID_PCM_MULAW || codecID == AV_CODEC_ID_PCM_ALAW)
		{
			got = false;
		}
	}
	return got;
}